{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import random\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchtext\n",
    "import pytorch_lightning as pl\n",
    "from transformers import (\n",
    "    DataProcessor,\n",
    "    InputExample,\n",
    "    BertTokenizer,\n",
    "    BertForSequenceClassification,\n",
    "    BertConfig,\n",
    "    glue_convert_examples_to_features,\n",
    ")\n",
    "\n",
    "def set_random_seed(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.backends.cudnn.deterministic = True\n",
    "\n",
    "set_random_seed(2020)\n",
    "device = torch.device('cuda:0' if torch.cuda.is_available else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>entity</th>\n",
       "      <th>offset</th>\n",
       "      <th>rawtext</th>\n",
       "      <th>type</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>小品</td>\n",
       "      <td>0</td>\n",
       "      <td>小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人</td>\n",
       "      <td>Other</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>战狼故事</td>\n",
       "      <td>3</td>\n",
       "      <td>小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人</td>\n",
       "      <td>Work</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>吴京</td>\n",
       "      <td>10</td>\n",
       "      <td>小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人</td>\n",
       "      <td>Person</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>障碍</td>\n",
       "      <td>16</td>\n",
       "      <td>小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人</td>\n",
       "      <td>Other</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>爱人</td>\n",
       "      <td>20</td>\n",
       "      <td>小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人</td>\n",
       "      <td>Other</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  entity  offset                         rawtext    type\n",
       "0     小品       0  小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人   Other\n",
       "1   战狼故事       3  小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人    Work\n",
       "2     吴京      10  小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人  Person\n",
       "3     障碍      16  小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人   Other\n",
       "4     爱人      20  小品《战狼故事》中，吴京突破重重障碍解救爱人，深情告白太感人   Other"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_path = Path('/home/bnu/projects/CCKS2020-Entity-Linking/data/csv')\n",
    "train_data = pd.read_csv(data_path/'train_type.csv', sep='\\t')\n",
    "train_data.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "53\n"
     ]
    }
   ],
   "source": [
    "lengths = [len(t) for t in train_data['rawtext']]\n",
    "print(max(lengths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Person',\n",
       " 'Work',\n",
       " 'Medicine',\n",
       " 'Game',\n",
       " 'Other',\n",
       " 'Organization',\n",
       " 'Location',\n",
       " 'Culture',\n",
       " 'Biological',\n",
       " 'VirtualThings',\n",
       " 'Natural&Geography',\n",
       " 'Website',\n",
       " 'Event',\n",
       " 'Brand',\n",
       " 'Food',\n",
       " 'Awards',\n",
       " 'Time&Calendar',\n",
       " 'Disease&Symptom',\n",
       " 'Software',\n",
       " 'Vehicle',\n",
       " 'Education',\n",
       " 'Constellation',\n",
       " 'Diagnosis&Treatment',\n",
       " 'Law&Regulation']"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pickle_path = Path('/home/bnu/projects/CCKS2020-Entity-Linking/data/pickle')\n",
    "idx_to_type = pd.read_pickle(pickle_path/'idx_to_type.pkl')\n",
    "idx_to_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ETProcessor(DataProcessor):\n",
    "\n",
    "    def get_train_examples(self, data_dir):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(data_dir / 'train_type.csv'),\n",
    "            set_type='train',\n",
    "        )\n",
    "\n",
    "    def get_dev_examples(self, data_dir):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(data_dir / 'valid_type.csv'),\n",
    "            set_type='valid',\n",
    "        )\n",
    "\n",
    "    def get_labels(self):\n",
    "        return idx_to_type\n",
    "\n",
    "    def _create_examples(self, lines, set_type):\n",
    "        examples = []\n",
    "        for i, line in enumerate(lines):\n",
    "            if i == 0:\n",
    "                continue\n",
    "            guid = f'{set_type}-{i}'\n",
    "            try:\n",
    "                text_a = line[0]\n",
    "                text_b = line[2]\n",
    "                label = line[3]  \n",
    "                examples.append(InputExample(\n",
    "                    guid=guid,\n",
    "                    text_a=text_a,\n",
    "                    text_b=text_b,\n",
    "                    label=label,\n",
    "                ))\n",
    "            except:\n",
    "                print(i)\n",
    "                print(line)\n",
    "        return examples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "InputExample(guid='train-11', text_a='动作', text_b='甄嬛传：安陵容怀孕时，雍正经常摸她的肚子，原来这动作大有深意', label='Other')\n",
      "Train: 266740\n",
      "Valid: 33074\n"
     ]
    }
   ],
   "source": [
    "data_path = Path('/home/bnu/projects/CCKS2020-Entity-Linking/data/csv')\n",
    "processor = ETProcessor()\n",
    "examples = processor.get_train_examples(data_path)\n",
    "print(examples[10])\n",
    "print('Train:', len(examples))\n",
    "examples = processor.get_dev_examples(data_path)\n",
    "print('Valid:', len(examples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_dataloaders(tokenizer, data_path):\n",
    "    def generate_dataloader_inner(examples):\n",
    "        features = glue_convert_examples_to_features(\n",
    "            examples,\n",
    "            tokenizer,\n",
    "            label_list=idx_to_type,\n",
    "            max_length=64,\n",
    "            output_mode='classification',\n",
    "            pad_on_left=False,\n",
    "            pad_token=tokenizer.pad_token_id,\n",
    "            pad_token_segment_id=0)\n",
    "\n",
    "        dataset = torch.utils.data.TensorDataset(\n",
    "            torch.LongTensor([f.input_ids for f in features]),\n",
    "            torch.LongTensor([f.attention_mask for f in features]),\n",
    "            torch.LongTensor([f.token_type_ids for f in features]),\n",
    "            torch.LongTensor([f.label for f in features])\n",
    "        )\n",
    "\n",
    "        sampler = torch.utils.data.RandomSampler(dataset)\n",
    "        dataloader = torch.utils.data.DataLoader(\n",
    "            dataset, sampler=sampler, batch_size=64\n",
    "        )\n",
    "        return dataloader\n",
    "\n",
    "    # 训练数据\n",
    "    train_examples = processor.get_train_examples(data_path)\n",
    "    print('Load Example Finish')\n",
    "    train_loader = generate_dataloader_inner(train_examples)\n",
    "    print('Generate DataLoader Finish')\n",
    "    \n",
    "    # 验证数据\n",
    "    valid_examples = processor.get_dev_examples(data_path)\n",
    "    print('Load Example Finish')\n",
    "    valid_loader = generate_dataloader_inner(valid_examples)\n",
    "    print('Generate DataLoader Finish')\n",
    "    \n",
    "    return train_loader, valid_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.tokenization_utils:Model name '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch' not found in model shortcut name list (bert-base-uncased, bert-large-uncased, bert-base-cased, bert-large-cased, bert-base-multilingual-uncased, bert-base-multilingual-cased, bert-base-chinese, bert-base-german-cased, bert-large-uncased-whole-word-masking, bert-large-cased-whole-word-masking, bert-large-uncased-whole-word-masking-finetuned-squad, bert-large-cased-whole-word-masking-finetuned-squad, bert-base-cased-finetuned-mrpc, bert-base-german-dbmdz-cased, bert-base-german-dbmdz-uncased, bert-base-finnish-cased-v1, bert-base-finnish-uncased-v1, bert-base-dutch-cased). Assuming '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch' is a path, a model identifier, or url to a directory containing tokenizer files.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/added_tokens.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/special_tokens_map.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:Didn't find file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/tokenizer_config.json. We won't load it.\n",
      "INFO:transformers.tokenization_utils:loading file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/vocab.txt\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.tokenization_utils:loading file None\n",
      "INFO:transformers.data.processors.glue:Writing example 0/266740\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: train-1\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 2207 1501 102 2207 1501 517 2773 4331 3125 752 518 704 8024 1426 776 4960 4788 7028 7028 7397 4809 6237 3131 4263 782 8024 3918 2658 1440 4635 1922 2697 782 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Other (id = 4)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: train-2\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 2773 4331 3125 752 102 2207 1501 517 2773 4331 3125 752 518 704 8024 1426 776 4960 4788 7028 7028 7397 4809 6237 3131 4263 782 8024 3918 2658 1440 4635 1922 2697 782 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Work (id = 1)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: train-3\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 1426 776 102 2207 1501 517 2773 4331 3125 752 518 704 8024 1426 776 4960 4788 7028 7028 7397 4809 6237 3131 4263 782 8024 3918 2658 1440 4635 1922 2697 782 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Person (id = 0)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: train-4\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 7397 4809 102 2207 1501 517 2773 4331 3125 752 518 704 8024 1426 776 4960 4788 7028 7028 7397 4809 6237 3131 4263 782 8024 3918 2658 1440 4635 1922 2697 782 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Other (id = 4)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: train-5\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 4263 782 102 2207 1501 517 2773 4331 3125 752 518 704 8024 1426 776 4960 4788 7028 7028 7397 4809 6237 3131 4263 782 8024 3918 2658 1440 4635 1922 2697 782 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Other (id = 4)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Load Example Finish\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.data.processors.glue:Writing example 10000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 20000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 30000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 40000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 50000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 60000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 70000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 80000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 90000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 100000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 110000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 120000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 130000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 140000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 150000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 160000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 170000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 180000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 190000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 200000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 210000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 220000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 230000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 240000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 250000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 260000/266740\n",
      "INFO:transformers.data.processors.glue:Writing example 0/33074\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: valid-1\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 1921 678 3766 3300 679 3141 4638 2155 2375 102 1921 678 3766 3300 679 3141 4638 2155 2375 118 446 9775 7437 704 2359 446 9775 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Work (id = 1)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: valid-2\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 446 9775 7437 704 2359 446 9775 102 1921 678 3766 3300 679 3141 4638 2155 2375 118 446 9775 7437 704 2359 446 9775 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Other (id = 4)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: valid-3\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 3719 1649 102 3719 1649 1322 2791 1139 4909 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Location (id = 6)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: valid-4\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 1322 2791 102 3719 1649 1322 2791 1139 4909 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Location (id = 6)\n",
      "INFO:transformers.data.processors.glue:*** Example ***\n",
      "INFO:transformers.data.processors.glue:guid: valid-5\n",
      "INFO:transformers.data.processors.glue:input_ids: 101 1139 4909 102 3719 1649 1322 2791 1139 4909 102 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:attention_mask: 1 1 1 1 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:token_type_ids: 0 0 0 0 1 1 1 1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      "INFO:transformers.data.processors.glue:label: Other (id = 4)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate DataLoader Finish\n",
      "Load Example Finish\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.data.processors.glue:Writing example 10000/33074\n",
      "INFO:transformers.data.processors.glue:Writing example 20000/33074\n",
      "INFO:transformers.data.processors.glue:Writing example 30000/33074\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generate DataLoader Finish\n"
     ]
    }
   ],
   "source": [
    "pretrained_path = '/media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch'\n",
    "tokenizer = BertTokenizer.from_pretrained(pretrained_path)\n",
    "train_loader, valid_loader = generate_dataloaders(tokenizer, data_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ETRoBERTaModel(pl.LightningModule):\n",
    "\n",
    "    def __init__(self, \n",
    "                 pretrained_path, \n",
    "                 train_loader, \n",
    "                 valid_loader):\n",
    "        super(ETRoBERTaModel, self).__init__()\n",
    "        self.train_loader = train_loader\n",
    "        self.valid_loader = valid_loader\n",
    "\n",
    "        config = BertConfig.from_json_file(pretrained_path+'/bert_config.json')\n",
    "        config.num_labels = len(idx_to_type)\n",
    "        print(config)\n",
    "        \n",
    "        # 预训练模型\n",
    "        self.ptm = BertForSequenceClassification.from_pretrained(\n",
    "            pretrained_path+'/pytorch_model.bin',\n",
    "            config=config,\n",
    "        )\n",
    "\n",
    "        # 损失函数\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, token_type_ids):\n",
    "        return self.ptm(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids\n",
    "        )[0]\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, label = batch\n",
    "        out = self(input_ids, attention_mask, token_type_ids)\n",
    "\n",
    "        loss = self.criterion(out, label)\n",
    "\n",
    "        _, pred = torch.max(out, dim=1)\n",
    "        acc = (pred == label).float().mean()\n",
    "\n",
    "        tensorboard_logs = {'train_loss': loss, 'train_acc': acc}\n",
    "        return {'loss': loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, label = batch\n",
    "        out = self(input_ids, attention_mask, token_type_ids)\n",
    "\n",
    "        loss = self.criterion(out, label)\n",
    "\n",
    "        _, pred = torch.max(out, dim=1)\n",
    "        acc = (pred == label).float().mean()\n",
    "\n",
    "        return {'val_loss': loss, 'val_acc': acc}\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
    "        val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()\n",
    "\n",
    "        tensorboard_logs = {'val_loss': val_loss, 'val_acc': val_acc}\n",
    "        return {'val_loss': val_loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam([p for p in self.parameters() if p.requires_grad], lr=2e-5, eps=1e-8)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.train_loader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.valid_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.modeling_utils:loading weights file /media/bnu/data/transformers-pretrained-model/chinese_roberta_wwm_ext_pytorch/pytorch_model.bin\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "BertConfig {\n",
      "  \"_num_labels\": 24,\n",
      "  \"architectures\": null,\n",
      "  \"attention_probs_dropout_prob\": 0.1,\n",
      "  \"bad_words_ids\": null,\n",
      "  \"bos_token_id\": null,\n",
      "  \"decoder_start_token_id\": null,\n",
      "  \"directionality\": \"bidi\",\n",
      "  \"do_sample\": false,\n",
      "  \"early_stopping\": false,\n",
      "  \"eos_token_id\": null,\n",
      "  \"finetuning_task\": null,\n",
      "  \"hidden_act\": \"gelu\",\n",
      "  \"hidden_dropout_prob\": 0.1,\n",
      "  \"hidden_size\": 768,\n",
      "  \"id2label\": {\n",
      "    \"0\": \"LABEL_0\",\n",
      "    \"1\": \"LABEL_1\",\n",
      "    \"2\": \"LABEL_2\",\n",
      "    \"3\": \"LABEL_3\",\n",
      "    \"4\": \"LABEL_4\",\n",
      "    \"5\": \"LABEL_5\",\n",
      "    \"6\": \"LABEL_6\",\n",
      "    \"7\": \"LABEL_7\",\n",
      "    \"8\": \"LABEL_8\",\n",
      "    \"9\": \"LABEL_9\",\n",
      "    \"10\": \"LABEL_10\",\n",
      "    \"11\": \"LABEL_11\",\n",
      "    \"12\": \"LABEL_12\",\n",
      "    \"13\": \"LABEL_13\",\n",
      "    \"14\": \"LABEL_14\",\n",
      "    \"15\": \"LABEL_15\",\n",
      "    \"16\": \"LABEL_16\",\n",
      "    \"17\": \"LABEL_17\",\n",
      "    \"18\": \"LABEL_18\",\n",
      "    \"19\": \"LABEL_19\",\n",
      "    \"20\": \"LABEL_20\",\n",
      "    \"21\": \"LABEL_21\",\n",
      "    \"22\": \"LABEL_22\",\n",
      "    \"23\": \"LABEL_23\"\n",
      "  },\n",
      "  \"initializer_range\": 0.02,\n",
      "  \"intermediate_size\": 3072,\n",
      "  \"is_decoder\": false,\n",
      "  \"is_encoder_decoder\": false,\n",
      "  \"label2id\": {\n",
      "    \"LABEL_0\": 0,\n",
      "    \"LABEL_1\": 1,\n",
      "    \"LABEL_10\": 10,\n",
      "    \"LABEL_11\": 11,\n",
      "    \"LABEL_12\": 12,\n",
      "    \"LABEL_13\": 13,\n",
      "    \"LABEL_14\": 14,\n",
      "    \"LABEL_15\": 15,\n",
      "    \"LABEL_16\": 16,\n",
      "    \"LABEL_17\": 17,\n",
      "    \"LABEL_18\": 18,\n",
      "    \"LABEL_19\": 19,\n",
      "    \"LABEL_2\": 2,\n",
      "    \"LABEL_20\": 20,\n",
      "    \"LABEL_21\": 21,\n",
      "    \"LABEL_22\": 22,\n",
      "    \"LABEL_23\": 23,\n",
      "    \"LABEL_3\": 3,\n",
      "    \"LABEL_4\": 4,\n",
      "    \"LABEL_5\": 5,\n",
      "    \"LABEL_6\": 6,\n",
      "    \"LABEL_7\": 7,\n",
      "    \"LABEL_8\": 8,\n",
      "    \"LABEL_9\": 9\n",
      "  },\n",
      "  \"layer_norm_eps\": 1e-12,\n",
      "  \"length_penalty\": 1.0,\n",
      "  \"max_length\": 20,\n",
      "  \"max_position_embeddings\": 512,\n",
      "  \"min_length\": 0,\n",
      "  \"model_type\": \"bert\",\n",
      "  \"no_repeat_ngram_size\": 0,\n",
      "  \"num_attention_heads\": 12,\n",
      "  \"num_beams\": 1,\n",
      "  \"num_hidden_layers\": 12,\n",
      "  \"num_return_sequences\": 1,\n",
      "  \"output_attentions\": false,\n",
      "  \"output_hidden_states\": false,\n",
      "  \"output_past\": true,\n",
      "  \"pad_token_id\": 0,\n",
      "  \"pooler_fc_size\": 768,\n",
      "  \"pooler_num_attention_heads\": 12,\n",
      "  \"pooler_num_fc_layers\": 3,\n",
      "  \"pooler_size_per_head\": 128,\n",
      "  \"pooler_type\": \"first_token_transform\",\n",
      "  \"prefix\": null,\n",
      "  \"pruned_heads\": {},\n",
      "  \"repetition_penalty\": 1.0,\n",
      "  \"task_specific_params\": null,\n",
      "  \"temperature\": 1.0,\n",
      "  \"top_k\": 50,\n",
      "  \"top_p\": 1.0,\n",
      "  \"torchscript\": false,\n",
      "  \"type_vocab_size\": 2,\n",
      "  \"use_bfloat16\": false,\n",
      "  \"vocab_size\": 21128\n",
      "}\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:transformers.modeling_utils:Weights of BertForSequenceClassification not initialized from pretrained model: ['classifier.weight', 'classifier.bias']\n",
      "INFO:transformers.modeling_utils:Weights from pretrained model not used in BertForSequenceClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.5866,  0.0903,  0.2801,  ..., -0.0650, -1.0424,  0.5130],\n",
       "        [ 0.5835,  0.2605,  0.1885,  ...,  0.0015, -0.8778,  0.4723],\n",
       "        [ 0.2553,  0.1412,  0.1660,  ..., -0.0530, -0.8117,  0.3629],\n",
       "        ...,\n",
       "        [ 0.4258, -0.1628,  0.3814,  ..., -0.1921, -0.8992,  0.4901],\n",
       "        [ 0.1804,  0.2798,  0.2134,  ..., -0.0246, -0.7629,  0.4898],\n",
       "        [ 0.2131, -0.0873,  0.1173,  ...,  0.1244, -0.7905,  0.4977]],\n",
       "       grad_fn=<AddmmBackward>)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = ETRoBERTaModel(pretrained_path, train_loader, valid_loader)\n",
    "batch = next(iter(train_loader))\n",
    "model(batch[0], batch[1], batch[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:lightning:GPU available: True, used: True\n",
      "INFO:lightning:VISIBLE GPUS: 0,1\n",
      "INFO:lightning:\n",
      "    | Name                                                 | Type                          | Params\n",
      "---------------------------------------------------------------------------------------------------\n",
      "0   | ptm                                                  | BertForSequenceClassification | 102 M \n",
      "1   | ptm.bert                                             | BertModel                     | 102 M \n",
      "2   | ptm.bert.embeddings                                  | BertEmbeddings                | 16 M  \n",
      "3   | ptm.bert.embeddings.word_embeddings                  | Embedding                     | 16 M  \n",
      "4   | ptm.bert.embeddings.position_embeddings              | Embedding                     | 393 K \n",
      "5   | ptm.bert.embeddings.token_type_embeddings            | Embedding                     | 1 K   \n",
      "6   | ptm.bert.embeddings.LayerNorm                        | LayerNorm                     | 1 K   \n",
      "7   | ptm.bert.embeddings.dropout                          | Dropout                       | 0     \n",
      "8   | ptm.bert.encoder                                     | BertEncoder                   | 85 M  \n",
      "9   | ptm.bert.encoder.layer                               | ModuleList                    | 85 M  \n",
      "10  | ptm.bert.encoder.layer.0                             | BertLayer                     | 7 M   \n",
      "11  | ptm.bert.encoder.layer.0.attention                   | BertAttention                 | 2 M   \n",
      "12  | ptm.bert.encoder.layer.0.attention.self              | BertSelfAttention             | 1 M   \n",
      "13  | ptm.bert.encoder.layer.0.attention.self.query        | Linear                        | 590 K \n",
      "14  | ptm.bert.encoder.layer.0.attention.self.key          | Linear                        | 590 K \n",
      "15  | ptm.bert.encoder.layer.0.attention.self.value        | Linear                        | 590 K \n",
      "16  | ptm.bert.encoder.layer.0.attention.self.dropout      | Dropout                       | 0     \n",
      "17  | ptm.bert.encoder.layer.0.attention.output            | BertSelfOutput                | 592 K \n",
      "18  | ptm.bert.encoder.layer.0.attention.output.dense      | Linear                        | 590 K \n",
      "19  | ptm.bert.encoder.layer.0.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "20  | ptm.bert.encoder.layer.0.attention.output.dropout    | Dropout                       | 0     \n",
      "21  | ptm.bert.encoder.layer.0.intermediate                | BertIntermediate              | 2 M   \n",
      "22  | ptm.bert.encoder.layer.0.intermediate.dense          | Linear                        | 2 M   \n",
      "23  | ptm.bert.encoder.layer.0.output                      | BertOutput                    | 2 M   \n",
      "24  | ptm.bert.encoder.layer.0.output.dense                | Linear                        | 2 M   \n",
      "25  | ptm.bert.encoder.layer.0.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "26  | ptm.bert.encoder.layer.0.output.dropout              | Dropout                       | 0     \n",
      "27  | ptm.bert.encoder.layer.1                             | BertLayer                     | 7 M   \n",
      "28  | ptm.bert.encoder.layer.1.attention                   | BertAttention                 | 2 M   \n",
      "29  | ptm.bert.encoder.layer.1.attention.self              | BertSelfAttention             | 1 M   \n",
      "30  | ptm.bert.encoder.layer.1.attention.self.query        | Linear                        | 590 K \n",
      "31  | ptm.bert.encoder.layer.1.attention.self.key          | Linear                        | 590 K \n",
      "32  | ptm.bert.encoder.layer.1.attention.self.value        | Linear                        | 590 K \n",
      "33  | ptm.bert.encoder.layer.1.attention.self.dropout      | Dropout                       | 0     \n",
      "34  | ptm.bert.encoder.layer.1.attention.output            | BertSelfOutput                | 592 K \n",
      "35  | ptm.bert.encoder.layer.1.attention.output.dense      | Linear                        | 590 K \n",
      "36  | ptm.bert.encoder.layer.1.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "37  | ptm.bert.encoder.layer.1.attention.output.dropout    | Dropout                       | 0     \n",
      "38  | ptm.bert.encoder.layer.1.intermediate                | BertIntermediate              | 2 M   \n",
      "39  | ptm.bert.encoder.layer.1.intermediate.dense          | Linear                        | 2 M   \n",
      "40  | ptm.bert.encoder.layer.1.output                      | BertOutput                    | 2 M   \n",
      "41  | ptm.bert.encoder.layer.1.output.dense                | Linear                        | 2 M   \n",
      "42  | ptm.bert.encoder.layer.1.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "43  | ptm.bert.encoder.layer.1.output.dropout              | Dropout                       | 0     \n",
      "44  | ptm.bert.encoder.layer.2                             | BertLayer                     | 7 M   \n",
      "45  | ptm.bert.encoder.layer.2.attention                   | BertAttention                 | 2 M   \n",
      "46  | ptm.bert.encoder.layer.2.attention.self              | BertSelfAttention             | 1 M   \n",
      "47  | ptm.bert.encoder.layer.2.attention.self.query        | Linear                        | 590 K \n",
      "48  | ptm.bert.encoder.layer.2.attention.self.key          | Linear                        | 590 K \n",
      "49  | ptm.bert.encoder.layer.2.attention.self.value        | Linear                        | 590 K \n",
      "50  | ptm.bert.encoder.layer.2.attention.self.dropout      | Dropout                       | 0     \n",
      "51  | ptm.bert.encoder.layer.2.attention.output            | BertSelfOutput                | 592 K \n",
      "52  | ptm.bert.encoder.layer.2.attention.output.dense      | Linear                        | 590 K \n",
      "53  | ptm.bert.encoder.layer.2.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "54  | ptm.bert.encoder.layer.2.attention.output.dropout    | Dropout                       | 0     \n",
      "55  | ptm.bert.encoder.layer.2.intermediate                | BertIntermediate              | 2 M   \n",
      "56  | ptm.bert.encoder.layer.2.intermediate.dense          | Linear                        | 2 M   \n",
      "57  | ptm.bert.encoder.layer.2.output                      | BertOutput                    | 2 M   \n",
      "58  | ptm.bert.encoder.layer.2.output.dense                | Linear                        | 2 M   \n",
      "59  | ptm.bert.encoder.layer.2.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "60  | ptm.bert.encoder.layer.2.output.dropout              | Dropout                       | 0     \n",
      "61  | ptm.bert.encoder.layer.3                             | BertLayer                     | 7 M   \n",
      "62  | ptm.bert.encoder.layer.3.attention                   | BertAttention                 | 2 M   \n",
      "63  | ptm.bert.encoder.layer.3.attention.self              | BertSelfAttention             | 1 M   \n",
      "64  | ptm.bert.encoder.layer.3.attention.self.query        | Linear                        | 590 K \n",
      "65  | ptm.bert.encoder.layer.3.attention.self.key          | Linear                        | 590 K \n",
      "66  | ptm.bert.encoder.layer.3.attention.self.value        | Linear                        | 590 K \n",
      "67  | ptm.bert.encoder.layer.3.attention.self.dropout      | Dropout                       | 0     \n",
      "68  | ptm.bert.encoder.layer.3.attention.output            | BertSelfOutput                | 592 K \n",
      "69  | ptm.bert.encoder.layer.3.attention.output.dense      | Linear                        | 590 K \n",
      "70  | ptm.bert.encoder.layer.3.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "71  | ptm.bert.encoder.layer.3.attention.output.dropout    | Dropout                       | 0     \n",
      "72  | ptm.bert.encoder.layer.3.intermediate                | BertIntermediate              | 2 M   \n",
      "73  | ptm.bert.encoder.layer.3.intermediate.dense          | Linear                        | 2 M   \n",
      "74  | ptm.bert.encoder.layer.3.output                      | BertOutput                    | 2 M   \n",
      "75  | ptm.bert.encoder.layer.3.output.dense                | Linear                        | 2 M   \n",
      "76  | ptm.bert.encoder.layer.3.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "77  | ptm.bert.encoder.layer.3.output.dropout              | Dropout                       | 0     \n",
      "78  | ptm.bert.encoder.layer.4                             | BertLayer                     | 7 M   \n",
      "79  | ptm.bert.encoder.layer.4.attention                   | BertAttention                 | 2 M   \n",
      "80  | ptm.bert.encoder.layer.4.attention.self              | BertSelfAttention             | 1 M   \n",
      "81  | ptm.bert.encoder.layer.4.attention.self.query        | Linear                        | 590 K \n",
      "82  | ptm.bert.encoder.layer.4.attention.self.key          | Linear                        | 590 K \n",
      "83  | ptm.bert.encoder.layer.4.attention.self.value        | Linear                        | 590 K \n",
      "84  | ptm.bert.encoder.layer.4.attention.self.dropout      | Dropout                       | 0     \n",
      "85  | ptm.bert.encoder.layer.4.attention.output            | BertSelfOutput                | 592 K \n",
      "86  | ptm.bert.encoder.layer.4.attention.output.dense      | Linear                        | 590 K \n",
      "87  | ptm.bert.encoder.layer.4.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "88  | ptm.bert.encoder.layer.4.attention.output.dropout    | Dropout                       | 0     \n",
      "89  | ptm.bert.encoder.layer.4.intermediate                | BertIntermediate              | 2 M   \n",
      "90  | ptm.bert.encoder.layer.4.intermediate.dense          | Linear                        | 2 M   \n",
      "91  | ptm.bert.encoder.layer.4.output                      | BertOutput                    | 2 M   \n",
      "92  | ptm.bert.encoder.layer.4.output.dense                | Linear                        | 2 M   \n",
      "93  | ptm.bert.encoder.layer.4.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "94  | ptm.bert.encoder.layer.4.output.dropout              | Dropout                       | 0     \n",
      "95  | ptm.bert.encoder.layer.5                             | BertLayer                     | 7 M   \n",
      "96  | ptm.bert.encoder.layer.5.attention                   | BertAttention                 | 2 M   \n",
      "97  | ptm.bert.encoder.layer.5.attention.self              | BertSelfAttention             | 1 M   \n",
      "98  | ptm.bert.encoder.layer.5.attention.self.query        | Linear                        | 590 K \n",
      "99  | ptm.bert.encoder.layer.5.attention.self.key          | Linear                        | 590 K \n",
      "100 | ptm.bert.encoder.layer.5.attention.self.value        | Linear                        | 590 K \n",
      "101 | ptm.bert.encoder.layer.5.attention.self.dropout      | Dropout                       | 0     \n",
      "102 | ptm.bert.encoder.layer.5.attention.output            | BertSelfOutput                | 592 K \n",
      "103 | ptm.bert.encoder.layer.5.attention.output.dense      | Linear                        | 590 K \n",
      "104 | ptm.bert.encoder.layer.5.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "105 | ptm.bert.encoder.layer.5.attention.output.dropout    | Dropout                       | 0     \n",
      "106 | ptm.bert.encoder.layer.5.intermediate                | BertIntermediate              | 2 M   \n",
      "107 | ptm.bert.encoder.layer.5.intermediate.dense          | Linear                        | 2 M   \n",
      "108 | ptm.bert.encoder.layer.5.output                      | BertOutput                    | 2 M   \n",
      "109 | ptm.bert.encoder.layer.5.output.dense                | Linear                        | 2 M   \n",
      "110 | ptm.bert.encoder.layer.5.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "111 | ptm.bert.encoder.layer.5.output.dropout              | Dropout                       | 0     \n",
      "112 | ptm.bert.encoder.layer.6                             | BertLayer                     | 7 M   \n",
      "113 | ptm.bert.encoder.layer.6.attention                   | BertAttention                 | 2 M   \n",
      "114 | ptm.bert.encoder.layer.6.attention.self              | BertSelfAttention             | 1 M   \n",
      "115 | ptm.bert.encoder.layer.6.attention.self.query        | Linear                        | 590 K \n",
      "116 | ptm.bert.encoder.layer.6.attention.self.key          | Linear                        | 590 K \n",
      "117 | ptm.bert.encoder.layer.6.attention.self.value        | Linear                        | 590 K \n",
      "118 | ptm.bert.encoder.layer.6.attention.self.dropout      | Dropout                       | 0     \n",
      "119 | ptm.bert.encoder.layer.6.attention.output            | BertSelfOutput                | 592 K \n",
      "120 | ptm.bert.encoder.layer.6.attention.output.dense      | Linear                        | 590 K \n",
      "121 | ptm.bert.encoder.layer.6.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "122 | ptm.bert.encoder.layer.6.attention.output.dropout    | Dropout                       | 0     \n",
      "123 | ptm.bert.encoder.layer.6.intermediate                | BertIntermediate              | 2 M   \n",
      "124 | ptm.bert.encoder.layer.6.intermediate.dense          | Linear                        | 2 M   \n",
      "125 | ptm.bert.encoder.layer.6.output                      | BertOutput                    | 2 M   \n",
      "126 | ptm.bert.encoder.layer.6.output.dense                | Linear                        | 2 M   \n",
      "127 | ptm.bert.encoder.layer.6.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "128 | ptm.bert.encoder.layer.6.output.dropout              | Dropout                       | 0     \n",
      "129 | ptm.bert.encoder.layer.7                             | BertLayer                     | 7 M   \n",
      "130 | ptm.bert.encoder.layer.7.attention                   | BertAttention                 | 2 M   \n",
      "131 | ptm.bert.encoder.layer.7.attention.self              | BertSelfAttention             | 1 M   \n",
      "132 | ptm.bert.encoder.layer.7.attention.self.query        | Linear                        | 590 K \n",
      "133 | ptm.bert.encoder.layer.7.attention.self.key          | Linear                        | 590 K \n",
      "134 | ptm.bert.encoder.layer.7.attention.self.value        | Linear                        | 590 K \n",
      "135 | ptm.bert.encoder.layer.7.attention.self.dropout      | Dropout                       | 0     \n",
      "136 | ptm.bert.encoder.layer.7.attention.output            | BertSelfOutput                | 592 K \n",
      "137 | ptm.bert.encoder.layer.7.attention.output.dense      | Linear                        | 590 K \n",
      "138 | ptm.bert.encoder.layer.7.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "139 | ptm.bert.encoder.layer.7.attention.output.dropout    | Dropout                       | 0     \n",
      "140 | ptm.bert.encoder.layer.7.intermediate                | BertIntermediate              | 2 M   \n",
      "141 | ptm.bert.encoder.layer.7.intermediate.dense          | Linear                        | 2 M   \n",
      "142 | ptm.bert.encoder.layer.7.output                      | BertOutput                    | 2 M   \n",
      "143 | ptm.bert.encoder.layer.7.output.dense                | Linear                        | 2 M   \n",
      "144 | ptm.bert.encoder.layer.7.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "145 | ptm.bert.encoder.layer.7.output.dropout              | Dropout                       | 0     \n",
      "146 | ptm.bert.encoder.layer.8                             | BertLayer                     | 7 M   \n",
      "147 | ptm.bert.encoder.layer.8.attention                   | BertAttention                 | 2 M   \n",
      "148 | ptm.bert.encoder.layer.8.attention.self              | BertSelfAttention             | 1 M   \n",
      "149 | ptm.bert.encoder.layer.8.attention.self.query        | Linear                        | 590 K \n",
      "150 | ptm.bert.encoder.layer.8.attention.self.key          | Linear                        | 590 K \n",
      "151 | ptm.bert.encoder.layer.8.attention.self.value        | Linear                        | 590 K \n",
      "152 | ptm.bert.encoder.layer.8.attention.self.dropout      | Dropout                       | 0     \n",
      "153 | ptm.bert.encoder.layer.8.attention.output            | BertSelfOutput                | 592 K \n",
      "154 | ptm.bert.encoder.layer.8.attention.output.dense      | Linear                        | 590 K \n",
      "155 | ptm.bert.encoder.layer.8.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "156 | ptm.bert.encoder.layer.8.attention.output.dropout    | Dropout                       | 0     \n",
      "157 | ptm.bert.encoder.layer.8.intermediate                | BertIntermediate              | 2 M   \n",
      "158 | ptm.bert.encoder.layer.8.intermediate.dense          | Linear                        | 2 M   \n",
      "159 | ptm.bert.encoder.layer.8.output                      | BertOutput                    | 2 M   \n",
      "160 | ptm.bert.encoder.layer.8.output.dense                | Linear                        | 2 M   \n",
      "161 | ptm.bert.encoder.layer.8.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "162 | ptm.bert.encoder.layer.8.output.dropout              | Dropout                       | 0     \n",
      "163 | ptm.bert.encoder.layer.9                             | BertLayer                     | 7 M   \n",
      "164 | ptm.bert.encoder.layer.9.attention                   | BertAttention                 | 2 M   \n",
      "165 | ptm.bert.encoder.layer.9.attention.self              | BertSelfAttention             | 1 M   \n",
      "166 | ptm.bert.encoder.layer.9.attention.self.query        | Linear                        | 590 K \n",
      "167 | ptm.bert.encoder.layer.9.attention.self.key          | Linear                        | 590 K \n",
      "168 | ptm.bert.encoder.layer.9.attention.self.value        | Linear                        | 590 K \n",
      "169 | ptm.bert.encoder.layer.9.attention.self.dropout      | Dropout                       | 0     \n",
      "170 | ptm.bert.encoder.layer.9.attention.output            | BertSelfOutput                | 592 K \n",
      "171 | ptm.bert.encoder.layer.9.attention.output.dense      | Linear                        | 590 K \n",
      "172 | ptm.bert.encoder.layer.9.attention.output.LayerNorm  | LayerNorm                     | 1 K   \n",
      "173 | ptm.bert.encoder.layer.9.attention.output.dropout    | Dropout                       | 0     \n",
      "174 | ptm.bert.encoder.layer.9.intermediate                | BertIntermediate              | 2 M   \n",
      "175 | ptm.bert.encoder.layer.9.intermediate.dense          | Linear                        | 2 M   \n",
      "176 | ptm.bert.encoder.layer.9.output                      | BertOutput                    | 2 M   \n",
      "177 | ptm.bert.encoder.layer.9.output.dense                | Linear                        | 2 M   \n",
      "178 | ptm.bert.encoder.layer.9.output.LayerNorm            | LayerNorm                     | 1 K   \n",
      "179 | ptm.bert.encoder.layer.9.output.dropout              | Dropout                       | 0     \n",
      "180 | ptm.bert.encoder.layer.10                            | BertLayer                     | 7 M   \n",
      "181 | ptm.bert.encoder.layer.10.attention                  | BertAttention                 | 2 M   \n",
      "182 | ptm.bert.encoder.layer.10.attention.self             | BertSelfAttention             | 1 M   \n",
      "183 | ptm.bert.encoder.layer.10.attention.self.query       | Linear                        | 590 K \n",
      "184 | ptm.bert.encoder.layer.10.attention.self.key         | Linear                        | 590 K \n",
      "185 | ptm.bert.encoder.layer.10.attention.self.value       | Linear                        | 590 K \n",
      "186 | ptm.bert.encoder.layer.10.attention.self.dropout     | Dropout                       | 0     \n",
      "187 | ptm.bert.encoder.layer.10.attention.output           | BertSelfOutput                | 592 K \n",
      "188 | ptm.bert.encoder.layer.10.attention.output.dense     | Linear                        | 590 K \n",
      "189 | ptm.bert.encoder.layer.10.attention.output.LayerNorm | LayerNorm                     | 1 K   \n",
      "190 | ptm.bert.encoder.layer.10.attention.output.dropout   | Dropout                       | 0     \n",
      "191 | ptm.bert.encoder.layer.10.intermediate               | BertIntermediate              | 2 M   \n",
      "192 | ptm.bert.encoder.layer.10.intermediate.dense         | Linear                        | 2 M   \n",
      "193 | ptm.bert.encoder.layer.10.output                     | BertOutput                    | 2 M   \n",
      "194 | ptm.bert.encoder.layer.10.output.dense               | Linear                        | 2 M   \n",
      "195 | ptm.bert.encoder.layer.10.output.LayerNorm           | LayerNorm                     | 1 K   \n",
      "196 | ptm.bert.encoder.layer.10.output.dropout             | Dropout                       | 0     \n",
      "197 | ptm.bert.encoder.layer.11                            | BertLayer                     | 7 M   \n",
      "198 | ptm.bert.encoder.layer.11.attention                  | BertAttention                 | 2 M   \n",
      "199 | ptm.bert.encoder.layer.11.attention.self             | BertSelfAttention             | 1 M   \n",
      "200 | ptm.bert.encoder.layer.11.attention.self.query       | Linear                        | 590 K \n",
      "201 | ptm.bert.encoder.layer.11.attention.self.key         | Linear                        | 590 K \n",
      "202 | ptm.bert.encoder.layer.11.attention.self.value       | Linear                        | 590 K \n",
      "203 | ptm.bert.encoder.layer.11.attention.self.dropout     | Dropout                       | 0     \n",
      "204 | ptm.bert.encoder.layer.11.attention.output           | BertSelfOutput                | 592 K \n",
      "205 | ptm.bert.encoder.layer.11.attention.output.dense     | Linear                        | 590 K \n",
      "206 | ptm.bert.encoder.layer.11.attention.output.LayerNorm | LayerNorm                     | 1 K   \n",
      "207 | ptm.bert.encoder.layer.11.attention.output.dropout   | Dropout                       | 0     \n",
      "208 | ptm.bert.encoder.layer.11.intermediate               | BertIntermediate              | 2 M   \n",
      "209 | ptm.bert.encoder.layer.11.intermediate.dense         | Linear                        | 2 M   \n",
      "210 | ptm.bert.encoder.layer.11.output                     | BertOutput                    | 2 M   \n",
      "211 | ptm.bert.encoder.layer.11.output.dense               | Linear                        | 2 M   \n",
      "212 | ptm.bert.encoder.layer.11.output.LayerNorm           | LayerNorm                     | 1 K   \n",
      "213 | ptm.bert.encoder.layer.11.output.dropout             | Dropout                       | 0     \n",
      "214 | ptm.bert.pooler                                      | BertPooler                    | 590 K \n",
      "215 | ptm.bert.pooler.dense                                | Linear                        | 590 K \n",
      "216 | ptm.bert.pooler.activation                           | Tanh                          | 0     \n",
      "217 | ptm.dropout                                          | Dropout                       | 0     \n",
      "218 | ptm.classifier                                       | Linear                        | 18 K  \n",
      "219 | criterion                                            | CrossEntropyLoss              | 0     \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bnu/miniconda3/envs/nlp/lib/python3.7/site-packages/pytorch_lightning/utilities/warnings.py:18: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validation sanity check', layout=Layout(flex='2'), max=5.…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4d22d860cf964166baea9bfaa88946e2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', layout=Layout(flex='2'), max=1.0), HTML(value='')), …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/bnu/miniconda3/envs/nlp/lib/python3.7/site-packages/pytorch_lightning/utilities/warnings.py:18: RuntimeWarning: Displayed epoch numbers in the progress bar start from \"1\" until v0.6.x, but will start from \"0\" in v0.8.0.\n",
      "  warnings.warn(*args, **kwargs)\n",
      "/home/bnu/miniconda3/envs/nlp/lib/python3.7/site-packages/pytorch_lightning/utilities/warnings.py:18: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n",
      "/home/bnu/miniconda3/envs/nlp/lib/python3.7/site-packages/torch/nn/parallel/_functions.py:61: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, description='Validating', layout=Layout(flex='2'), max=517.0, style=Pr…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer = pl.Trainer(\n",
    "    max_epochs=1,\n",
    "    val_check_interval=0.1,\n",
    "    gpus=2,\n",
    "    distributed_backend='dp',\n",
    "    default_save_path='/media/bnu/data/pytorch-lightning-checkpoints/ETRoBERTaModel',\n",
    ")\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_path = Path('/home/bnu/projects/CCKS2020-Entity-Linking/ckpt/')\n",
    "trainer.save_checkpoint(ckpt_path/'ET-RoBERTa-64-0420.ckpt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
